import pandas as pd
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn
from sklearn.preprocessing import MinMaxScaler

class SensorDataset(Dataset):
    def __init__(self, list_file, csv_file, seq_length, label_length=20, step=1, device='cpu'):
        """
        Args:
            list_file (str): Path to the list.txt file containing sensor headers.
            csv_file (str): Path to the CSV file containing sensor data.
            seq_length (int): Length of each input sequence (seq_x).
            label_length (int): Length of each label sequence (seq_y).
            step (int): Step size for creating sequences.
            device (str): Device to store the data ('cpu' or 'cuda').
        """
        # 读取 list.txt，获取传感器名称
        with open(list_file, 'r') as f:
            self.headers = f.read().splitlines()
        
        # 读取CSV文件
        self.df = pd.read_csv(csv_file)
        
        # 检查所有传感器是否存在于CSV中
        missing_headers = [header for header in self.headers if header not in self.df.columns]
        if missing_headers:
            raise ValueError(f"以下传感器在CSV中未找到: {missing_headers}")
        
        # 提取传感器数据
        self.sensor_data = self.df[self.headers].values  # [num_samples, num_sensors]
        
        # 处理缺失值（这里选择填充为0，可以根据需要调整）
        if np.isnan(self.sensor_data).any():
            self.sensor_data = np.nan_to_num(self.sensor_data, nan=0.0)
        
        self.num_sensors = len(self.headers)
        self.seq_length = seq_length
        self.label_length = label_length
        self.step = step
        self.device = device

        # 生成序列
        self.sequences = []
        for i in range(0, self.sensor_data.shape[0] - seq_length - label_length + 1, step):
            seq_x = self.sensor_data[i:i + seq_length]  # [seq_length, num_sensors]
            seq_y = self.sensor_data[i + seq_length:i + seq_length + label_length]  # [label_length, num_sensors]
            self.sequences.append((seq_x, seq_y))

        self.sequences = np.array(self.sequences)  # [num_sequences, 2, seq_length or label_length, num_sensors]
        self.num_sequences = len(self.sequences)

    def __len__(self):
        return self.num_sequences

    def __getitem__(self, idx):
        # 返回一个序列及其标签
        seq_x, seq_y = self.sequences[idx]  # seq_x: [seq_length, num_sensors], seq_y: [label_length, num_sensors]
        seq_x = torch.tensor(seq_x, dtype=torch.float32, device=self.device)
        seq_y = torch.tensor(seq_y, dtype=torch.float32, device=self.device)
        return seq_x, seq_y

def collate_fn(batch):
    """
    Collate function to combine multiple (seq_x, seq_y) pairs into batches.
    Args:
        batch (list): List of (seq_x, seq_y) tuples.
    Returns:
        batch_x (torch.Tensor): [batch_size, seq_length, num_sensors]
        batch_y (torch.Tensor): [batch_size, label_length, num_sensors]
    """
    batch_x, batch_y = zip(*batch)
    batch_x = torch.stack(batch_x)  # [batch_size, seq_length, num_sensors]
    batch_y = torch.stack(batch_y)  # [batch_size, label_length, num_sensors]
    return batch_x, batch_y
